from PIL import Image
import numpy as np


def convert_to_grayscale(img):
    img_array = np.array(img)
    grayscale_array = np.zeros((img.height, img.width), dtype=np.uint8)

    for i in range(img.height):
        for j in range(img.width):
            sum = 0
            for num in img_array[i, j]:
                sum = sum + num
            grayscale_array[i, j] = sum / 3

    grayscale_img = Image.fromarray(grayscale_array, 'L')
    return grayscale_img, grayscale_array


def convert_to_binary(img):
    if img.mode == 'RGB':
        _, grayscale_array = convert_to_grayscale(img)
    else:
        grayscale_array = np.array(img)

    binary_array = np.zeros((img.height, img.width), dtype=np.uint8)

    for i in range(img.height):
        for j in range(img.width):
            binary_array[i, j] = 0 if grayscale_array[i, j] < 128 else 255

    binary_img = Image.fromarray(binary_array, 'L')  # if this won't work, set mode='L'
    binary_img.mode = '1'
    return binary_img


def get_hist_arr(h, w, img_arr):
    hist_arr = [0] * 256
    for i in range(h):
        for j in range(w):
            hist_arr[img_arr[i][j]] += 1
    return hist_arr


def get_cum_sum(hist_arr):
    hist_arr = iter(hist_arr)
    cum_sum = [next(hist_arr)]
    for i in hist_arr:
        cum_sum.append(cum_sum[-1] + i)
    return np.array(cum_sum)  # make it numpy for further use


def equalize(h, w, img_arr, hist_arr):
    cum_sum = get_cum_sum(hist_arr)

    # conform cumulative sum to 0-255
    # formula: s_k = sum(j=0, k) (n_j/N)
    nj = (cum_sum - cum_sum.min()) * 255
    N = cum_sum.max() - cum_sum.min()
    cum_sum = nj / N
    cum_sum = cum_sum.astype('uint8')  # cast it back to uint8

    # get the value form cum_sum for every index in img_arr,
    # and set that as img_arr
    new_arr = np.zeros((h, w), dtype='uint8')
    for i in range(h):
        for j in range(w):
            new_arr[i][j] = cum_sum[img_arr[i][j]]
    return new_arr


def clamp(num_list):
    ret = []
    for item in num_list:
        if item < 0:
            item = 0
        elif item > 255:
            item = 255
        ret.append(round(item))
    return ret


def rgb2ycbcr(h, w, img_arr):
    """
    Mathematically,
    [Y, Cb, Cr].T = [0, 128, 128].T + k1 * [R, G, B].T;
    """
    # slice the channels
    r = img_arr[:, :, 0]
    g = img_arr[:, :, 1]
    b = img_arr[:, :, 2]

    # convert to YCbCr color space first
    k1 = np.mat([[0.299, 0.587, 0.114],
                 [-0.169, -0.331, 0.5],
                 [0.5, -0.419, -0.081]])
    b1 = np.mat([0, 128, 128])

    y = np.zeros((h, w), dtype='uint8')
    cb = np.zeros((h, w), dtype='uint8')
    cr = np.zeros((h, w), dtype='uint8')
    for i in range(h):
        for j in range(w):
            y[i][j], cb[i][j], cr[i][j] = clamp(np.array(b1.T + k1 * np.mat([r[i][j], g[i][j], b[i][j]]).T)[:, 0])
    return y, cb, cr


def ycbcr2rgb(h, w, y, cb, cr):
    """
    Mathematically,
    [R, G, B].T = k2 * [Y, (Cb - 128), (Cr - 128)].T
    """
    k2 = np.mat([[1, 0, 1.4],
                 [1, -0.343, -0.711],
                 [1, 1.765, 0]])
    r = np.zeros((h, w), dtype='uint8')
    g = np.zeros((h, w), dtype='uint8')
    b = np.zeros((h, w), dtype='uint8')
    # convert back to rgb
    for i in range(h):
        for j in range(w):
            r[i][j], g[i][j], b[i][j] = clamp(
                np.array(k2 * np.mat([y[i][j], cb[i][j] - 128, cr[i][j] - 128]).T)[:, 0])
    return r, g, b


def rgb2array(h, w, r, g, b):
    # form new array
    new_arr = np.zeros((h, w, 3), dtype='uint8')
    for i in range(h):
        for j in range(w):
            new_arr[i][j][0] = r[i][j]
            new_arr[i][j][1] = g[i][j]
            new_arr[i][j][2] = b[i][j]
    return new_arr


def hist_equalize(img):
    h = img.height
    w = img.width
    img_arr = np.array(img)
    if img.mode == 'RGB':
        """  # to let pillow do it
        img_arr = np.array(img.convert('YCbCr'))
        print(img_arr)
        y = img_arr[:, :, 0]
        hist_arr = get_hist_arr(h, w, y)
        y = equalize(h, w, y, hist_arr)
        for i in range(h):
            for j in range(w):
                img_arr[i][j][0] = y[i][j]
        new_img = Image.fromarray(img_arr, mode='YCbCr')
        return new_img
        """
        y, cb, cr = rgb2ycbcr(h, w, img_arr)

        # equalize channel y
        hist_arr = get_hist_arr(h, w, y)
        y = equalize(h, w, y, hist_arr)

        r, g, b = ycbcr2rgb(h, w, y, cb, cr)
        new_arr = rgb2array(h, w, r, g, b)

        # write new image
        return Image.fromarray(new_arr, mode='RGB')

    else:  # if image is grayscale
        hist_arr = get_hist_arr(h, w, img_arr)
        new_arr = equalize(h, w, img_arr, hist_arr)
        return Image.fromarray(new_arr, mode='L')


# Returns index of element closest to target in arr[]
def find_closest(arr, n, target, last_j):
    min_diff = abs(target - arr[last_j])
    min_index = last_j
    for i in range(last_j + 1, n):
        diff = abs(arr[i] - target)
        if diff < min_diff:
            min_diff = diff
            min_index = i
        elif diff > min_diff:
            break
    return min_index


def regulate(h, w, img_arr, target_hist):
    cum_sum_1 = get_cum_sum(get_hist_arr(img_arr))
    cum_den_1 = np.zeros(256, dtype='f')  # this is necessary
    cum_den_1 = cum_sum_1 / cum_sum_1[255]

    cum_sum_2 = get_cum_sum(get_hist_arr(target_hist))
    cum_den_2 = np.zeros(256, dtype='f')  # this is necessary
    cum_den_2 = cum_sum_2 / cum_sum_2[255]

    # define mapping
    old = [i for i in range(256)]
    new = [0] * 256

    # find mapping
    j = 0
    for i in range(256):
        j = find_closest(cum_den_2, 156, cum_den_1[i], j)
        new[i] = j
    mapping = zip(old, new)
    mapping = dict(mapping)

    new_arr = np.zeros((h, w), dtype='uint8')
    for i in range(h):
        for j in range(w):
            new_arr[i][j] = mapping[img_arr[i][j]]
    return new_arr


def hist_regulate_with_sample(img, channel, target_hist):
    img_arr = np.array(img)
    h = img.height
    w = img.width

    if channel == 'RGB' and img.mode == 'RGB':  # consider using pillow's convert if this is too slow
        y, cb, cr = rgb2ycbcr(h, w, img_arr)
        new_y = regulate(h, w, y, target_hist)
        r, g, b = ycbcr2rgb(h, w, new_y, cb, cr)
        new_arr = rgb2array(h, w, r, g, b)
        return Image.fromarray(new_arr, mode='RGB')
    elif img.mode == 'L':
        new_arr = regulate(h, w, img_arr, target_hist)
        return Image.fromarray(new_arr)
    else:
        r = img_arr[:, :, 0]
        g = img_arr[:, :, 1]
        b = img_arr[:, :, 2]
        if channel == 'Red':
            new_r = regulate(h, w, r, target_hist)
            new_arr = rgb2array(h, w, new_r, g, b)
        elif channel == 'Green':
            new_g = regulate(h, w, g, target_hist)
            new_arr = rgb2array(h, w, r, new_g, b)
        else:
            new_b = regulate(h, w, b, target_hist)
            new_arr = rgb2array(h, w, r, g, new_b)
        return Image.fromarray(new_arr, mode='RGB')


def replace(h, w, img_arr, target):
    new_arr = np.zeros((h, w), dtype='uint8')
    for i in range(h):
        for j in range(w):
            new_arr[i][j] = target[img_arr[i][j]]
    return new_arr


def hist_regulate(img, channel, target):
    img_arr = np.array(img)
    h = img.height
    w = img.width

    if channel == 'RGB' and img.mode == 'RGB':  # consider using pillow's convert if this is too slow
        y, cb, cr = rgb2ycbcr(h, w, img_arr)
        new_y = replace(h, w, y, target)
        r, g, b = ycbcr2rgb(h, w, new_y, cb, cr)
        new_arr = rgb2array(h, w, r, g, b)
        return Image.fromarray(new_arr, mode='RGB')
    elif img.mode == 'L':
        new_arr = replace(h, w, img_arr, target)
        return Image.fromarray(new_arr)
    else:
        r = img_arr[:, :, 0]
        g = img_arr[:, :, 1]
        b = img_arr[:, :, 2]
        if channel == 'Red':
            new_r = replace(h, w, r, target)
            new_arr = rgb2array(h, w, new_r, g, b)
        elif channel == 'Green':
            new_g = replace(h, w, g, target)
            new_arr = rgb2array(h, w, r, new_g, b)
        else:
            new_b = replace(h, w, b, target)
            new_arr = rgb2array(h, w, r, g, new_b)
        return Image.fromarray(new_arr, mode='RGB')


def convolution(kernel):  # TODO: finish this
    pass
